package com.wfm.limitalgor;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicLong;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

@Slf4j
@Data
public class TokenBuketRateLimiter {
    // 上一次令牌发放时间
    public  long lastTime = 0;
    // 桶的容量
    public  int capacity = 1;
    // 令牌生成速度 /s
    public  int rate = 1;
    // 当前令牌数量
    public  AtomicLong tokens = new AtomicLong(0);

    public TokenBuketRateLimiter(int capacity,int rate) {
        this.capacity=capacity;
        this.rate=rate;
    }

    //返回值说明：
    // false 获取失败
    // true 获取成功
    public  synchronized boolean tryAcquire(int applyCount) {
        long now = System.currentTimeMillis();
        //时间间隔,单位为 ms
        long gap = now - lastTime;
        //场景三：当前请求和上次请求，在同一个时间区间
        //当前时间，在时间区间之内
        //以时间区间为计算维度，同一个区间，没有必要重复去计算令牌数量
        if (lastTime != 0 && gap < 1000) {
            if (tokens.get() < applyCount) {
                // 若拿不到令牌,则拒绝
                log.info("tryAcquire fail");
                return false;
            } else {
                // 还有令牌，领取令牌
                tokens.getAndAdd(-applyCount);
                log.info("tryAcquire success");
                return true;
            }
        }
        if (lastTime == 0) {
            gap = 1000;
        }
        //计算时间段内的令牌数
        int generatorTokens = (int) (gap * rate / 1000);
        int allTokens = Math.toIntExact(tokens.get() + generatorTokens);
        // 当前令牌数
        tokens.set(Math.min(capacity, allTokens));
        lastTime = now;
        if (tokens.get() < applyCount) {
            // 若拿不到令牌,则拒绝
            log.info("tryAcquire fail");
            return false;
        } else {
            // 还有令牌，领取令牌
            tokens.getAndAdd(-applyCount);
            log.info("tryAcquire success");
            return true;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        test2();
    }

    private static void test1() throws InterruptedException {
        TokenBuketRateLimiter limiter = new TokenBuketRateLimiter(3, 3);
        for (int i = 0; i < 4; i++) {
            System.out.println("================="+i+"=============");
            switch (i){
                case 0:
                    for (int j = 0; j < 3; j++) {
                        limiter.tryAcquire(1);
                    }
                    break;
                case 1:
                    for (int j = 0; j < 4; j++) {
                        limiter.tryAcquire(1);
                    }
                    break;
                case 2:
                    for (int j = 0; j < 5; j++) {
                        limiter.tryAcquire(1);
                    }
                    break;
                case 3:
                    for (int j = 0; j < 2; j++) {
                        limiter.tryAcquire(1);
                    }
                    break;
                default:
                    break;
            }

            Thread.sleep(1000);

        }
    }

    /**
     * 总共4个0.5s，就是2s
     * 第一个0.5s：不请求
     * 第二个0.5s: 请求3个
     * 第三个0.5s：不请求
     * 第四个0.5s：请求三个
     * 0__1__2__3
     * 在1~3这两个点，分别请求3，发现1s内请求成功了6次，复合令牌桶应对突发流量的功能
     *
     */
    private static void test2() throws InterruptedException {
        TokenBuketRateLimiter limiter = new TokenBuketRateLimiter(6, 3);
        for (int i = 0; i < 6; i++) {
            System.out.println("================="+i+"=============");
            switch (i){
                case 0:
                    break;
                case 1:
                    for (int j = 0; j < 3; j++) {
                        limiter.tryAcquire(1);
                    }
                    break;
                case 5:
                    for (int j = 0; j < 6; j++) {
                        limiter.tryAcquire(1);
                    }
                    break;
                default:
                    break;
            }

            Thread.sleep(500);

        }
    }

}
